{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5c0122e65c053f38",
   "metadata": {},
   "source": [
    "# Inbound Auth\n",
    "\n",
    "AgentCore Identity lets you validate inbound access (Inbound Auth) for users and applications calling agents or tools in an AgentCore Runtime or validate access to AgentCore Gateway targets. It also provide secure outbound access (Outbound Auth) from an agent to external services or a Gateway target. It integrates with your existing identity providers (such as Amazon Cognito) while enforcing permission boundaries for agents acting independently or on behalf of users (via OAuth).\n",
    "\n",
    "Inbound Auth validates callers attempting to invoke agents or tools, whether they're hosted in AgentCore Runtime, AgentCore Gateway , or in other environments. Inbound Auth works with IAM (SigV4 credentials) or with OAuth authorization.\n",
    "\n",
    "By default, Amazon Bedrock AgentCore uses IAM credentials, meaning user requests to the agent are authenticated with the user's IAM credentials. If you use OAuth, you will need to specify the following when configuring your AgentCore Runtime resources or AgentCore Gateway endpoints:\n",
    "\n",
    "- OAuth discovery server Url — A string that must match the pattern ^.+/\\.well-known/openid-configuration$ for OpenID Connect discovery URLs\n",
    "\n",
    "- Allowed audiences — List of allowed audiences for JWT tokens\n",
    "\n",
    "- Allowed clients — List of allowed client identifiers\n",
    "\n",
    "If you use the AgentCore CLI, you can specify the type of authorization (and OAuth discovery server) for an AgentCore Runtime when you use the **configure** command. You can also use the CreateAgentRuntime operation and Amazon Bedrock AgentCore console. If you are creating a Gateway, you use the CreateGateway operation, or the console.\n",
    "\n",
    "Before the user can use the agent, the client application must have the user authenticate with the OAuth authorizer. Your client receives a bearer token which it then passes to the agent in an invocation request. Upon receipt the agent validates the token with the authorization server before allowing access.\n",
    "\n",
    "\n",
    "## Overview\n",
    "\n",
    "In this tutorial we will modify the agent you deployed in 01-AgentCore-runtime and configure it for Inbound Auth using Cognito as the Identity provider. You will set up a Cognito User pool with one user and an app client. You will learn how to host your existing agent, using Amazon Bedrock AgentCore Runtime with Inbound Auth using the Cognito user pool. \n",
    "\n",
    "### Tutorial Architecture\n",
    "\n",
    "<div style=\"text-align:center\">\n",
    "    <img src=\"images/inbound_auth_cognito.png\" width=\"90%\"/>\n",
    "</div>\n",
    "\n",
    "### Tutorial Details\n",
    "\n",
    "\n",
    "| Information         | Details                                                                          |\n",
    "|:--------------------|:---------------------------------------------------------------------------------|\n",
    "| Tutorial type       | Conversational                                                                   |\n",
    "| Agent type          | Single                                                                           |\n",
    "| Agentic Framework   | Strands Agents                                                                   |\n",
    "| LLM model           | Anthropic Claude Sonnet 4                                                        |\n",
    "| Tutorial components | Hosting agent on AgentCore Runtime. Using Strands Agent and Amazon Bedrock Model |\n",
    "| Tutorial vertical   | Cross-vertical                                                                   |\n",
    "| Example complexity  | Easy                                                                             |\n",
    "| Inbound Auth        | Cognito                                                                          |\n",
    "| SDK used            | Amazon BedrockAgentCore Python SDK and boto3                                     |\n",
    "\n",
    "\n",
    "\n",
    "### Tutorial Key Features\n",
    "\n",
    "* Hosting Agents on Amazon Bedrock AgentCore Runtime with Inbound Auth using Amazon Cognito\n",
    "* Using Amazon Bedrock models\n",
    "* Using Strands Agents\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a676f58ecf52b42",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "To execute this tutorial you will need:\n",
    "* Python 3.10+\n",
    "* AWS credentials\n",
    "* Amazon Bedrock AgentCore SDK\n",
    "* Strands Agents\n",
    "* Docker running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {
    "jupyter": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "!uv add -r requirements.txt --active"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46ac3dfb",
   "metadata": {},
   "source": [
    "### Provision a Cognito User Pool\n",
    "\n",
    "Lets provision a Cognito Userpool with an App client and one test user. Note down the 1/Cognito Discovery url and 2/the Cognito app client id. We will use it to configure our agent for Inbound Auth with Cognito."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8af4d14",
   "metadata": {},
   "outputs": [],
   "source": [
    "!chmod +x setup_cognito_user_pool.sh\n",
    "!sh setup_cognito_user_pool.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "932110e6-fca6-47b6-b7c5-c4714a866a80",
   "metadata": {},
   "source": [
    "## Preparing your agent for deployment on AgentCore Runtime\n",
    "\n",
    "### Strands Agents with Amazon Bedrock model\n",
    "Let's start with our Strands Agent we created in the 01-AgentCore-runtime tutorial and configure it with Inbound Auth that uses Amazon Cognito as the Identity Provider."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b845b32-a03e-45c2-a2f0-2afba8069f47",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile strands_claude.py\n",
    "from strands import Agent, tool\n",
    "from strands_tools import calculator # Import the calculator tool\n",
    "import argparse\n",
    "import json\n",
    "from bedrock_agentcore.runtime import BedrockAgentCoreApp\n",
    "from strands.models import BedrockModel\n",
    "\n",
    "app = BedrockAgentCoreApp()\n",
    "\n",
    "# Create a custom tool \n",
    "@tool\n",
    "def weather():\n",
    "    \"\"\" Get weather \"\"\" # Dummy implementation\n",
    "    return \"sunny\"\n",
    "\n",
    "\n",
    "model_id = \"us.anthropic.claude-sonnet-4-20250514-v1:0\"\n",
    "model = BedrockModel(\n",
    "    model_id=model_id,\n",
    ")\n",
    "agent = Agent(\n",
    "    model=model,\n",
    "    tools=[calculator, weather],\n",
    "    system_prompt=\"You're a helpful assistant. You can do simple math calculation, and tell the weather.\"\n",
    ")\n",
    "\n",
    "@app.entrypoint\n",
    "def strands_agent_bedrock(payload):\n",
    "    \"\"\"\n",
    "    Invoke the agent with a payload\n",
    "    \"\"\"\n",
    "    user_input = payload.get(\"prompt\")\n",
    "    print(\"User input:\", user_input)\n",
    "    response = agent(user_input)\n",
    "    return response.message['content'][0]['text']\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    app.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6820ca8f-a8a8-4f34-b4ef-b6dad3776261",
   "metadata": {},
   "source": [
    "## Deploying the agent to AgentCore Runtime\n",
    "\n",
    "The `CreateAgentRuntime` operation supports comprehensive configuration options, letting you specify container images, environment variables and encryption settings. You can also configure protocol settings (HTTP, MCP) and authorization mechanisms to control how your clients communicate with the agent. \n",
    "\n",
    "**Note:** Operations best practice is to package code as container and push to ECR using CI/CD pipelines and IaC\n",
    "\n",
    "In this tutorial can will the Amazon Bedrock AgentCode Python SDK to easily package your artifacts and deploy them to AgentCore runtime."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0861401-a111-4ade-9e02-50f52fdfa9b1",
   "metadata": {},
   "source": [
    "### Creating runtime role\n",
    "\n",
    "Before starting, let's create an IAM role for our AgentCore Runtime. We will do so using the utils function pre-developed for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54dd2fdf-985c-4a70-8b87-071783a209de",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "\n",
    "# Get the current notebook's directory\n",
    "current_dir = os.path.dirname(os.path.abspath('__file__' if '__file__' in globals() else '.'))\n",
    "\n",
    "# Navigate up to the utils.py location\n",
    "utils_dir = os.path.join(current_dir, '..')\n",
    "utils_dir = os.path.abspath(utils_dir)\n",
    "\n",
    "# Add to sys.path\n",
    "sys.path.insert(0, utils_dir)\n",
    "\n",
    "from utils import create_agentcore_role\n",
    "\n",
    "agent_name=\"strands_claude\"\n",
    "agentcore_iam_role = create_agentcore_role(agent_name=agent_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8855aceb-b79f-4aaa-b16f-8577c059816a",
   "metadata": {},
   "source": [
    "### Configure AgentCore Runtime deployment\n",
    "\n",
    "Next we will use our starter toolkit to configure the AgentCore Runtime deployment with an entrypoint, the execution role we just created and a requirements file. We will also configure the starter kit to auto create the Amazon ECR repository on launch.\n",
    "\n",
    "During the configure step, your docker file will be generated based on your application code\n",
    "\n",
    "**Important** - Update the Cognito Discovery url and the Cognito App client id from the previous steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e79eba2-ca59-463f-9ebf-56e362d7ae66",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bedrock_agentcore_starter_toolkit import Runtime\n",
    "from boto3.session import Session\n",
    "boto_session = Session()\n",
    "region = boto_session.region_name\n",
    "region\n",
    "\n",
    "#Update the Cognito Discovery url below. You can get the discovery url from the \"Provision a Cognito User Pool\" section\n",
    "discovery_url = '<your-cognito-user-pool-discovery-url>' \n",
    "\n",
    "#Update the Cognito App Client id below. You can get the Cognito App client from the \"Provision a Cognito User Pool\" section\n",
    "client_id = '<your-cognito-app-client-id>' \n",
    "\n",
    "agentcore_runtime = Runtime()\n",
    "\n",
    "response = agentcore_runtime.configure(\n",
    "    entrypoint=\"strands_claude.py\",\n",
    "    execution_role=agentcore_iam_role['Role']['Arn'],\n",
    "    auto_create_ecr=True,\n",
    "    requirements_file=\"requirements.txt\",\n",
    "    region=region,\n",
    "    agent_name=agent_name+\"4\",\n",
    "    authorizer_configuration={\n",
    "        \"customJWTAuthorizer\": {\n",
    "            \"discoveryUrl\": discovery_url,\n",
    "            \"allowedClients\": [client_id]\n",
    "        }\n",
    "    }\n",
    ")\n",
    "response\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3d55172-602d-43df-9878-414f23285984",
   "metadata": {},
   "source": [
    "## Review the AgentCore configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2d71fbe-c3a9-4b4c-87f4-30b6137ab5f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat .bedrock_agentcore.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e1b84cc-798e-472c-ac0b-2c315f4b704d",
   "metadata": {},
   "source": [
    "### Launching agent to AgentCore Runtime\n",
    "\n",
    "Now that we've got a docker file, let's launch the agent to the AgentCore Runtime. This will create the Amazon ECR repository and the AgentCore Runtime\n",
    "\n",
    "<div style=\"text-align:left\">\n",
    "    <img src=\"images/launch.png\" width=\"75%\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17a32ab8-7701-4900-8055-e24364bdf35c",
   "metadata": {},
   "outputs": [],
   "source": [
    "launch_result = agentcore_runtime.launch()\n",
    "launch_result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0ae9c09-09db-4a76-871a-92eacd96b9c3",
   "metadata": {},
   "source": [
    "### Checking for the AgentCore Runtime Status\n",
    "Now that we've deployed the AgentCore Runtime, let's check for it's deployment status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afa6ac09-9adb-4846-9fc1-4d12aeb74853",
   "metadata": {},
   "outputs": [],
   "source": [
    "status_response = agentcore_runtime.status()\n",
    "status = status_response.endpoint['status']\n",
    "end_status = ['READY', 'CREATE_FAILED', 'DELETE_FAILED', 'UPDATE_FAILED']\n",
    "while status not in end_status:\n",
    "    time.sleep(10)\n",
    "    status_response = agentcore_runtime.status()\n",
    "    status = status_response.endpoint['status']\n",
    "    print(status)\n",
    "status"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7f89c56-918a-4cab-beaa-c7ac43a2ba29",
   "metadata": {},
   "source": [
    "### Invoking AgentCore Runtime without authorization\n",
    "\n",
    "Finally, we can invoke our AgentCore Runtime with a payload. Try running the following cell and you will see an error that says **\"AccessDeniedException: An error occurred (AccessDeniedException) when calling the InvokeAgentRuntime operation: Agent is configured for a different authorization token type\".**\n",
    "\n",
    "<div style=\"text-align:left\">\n",
    "    <img src=\"images/invoke.png\" width=75%\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d909e42-e1a0-407f-84c2-3d16cc889cd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "invoke_response = agentcore_runtime.invoke({\"prompt\": \"How is the weather now?\"})\n",
    "invoke_response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18d0d295-8279-499a-9a1b-3310086cf493",
   "metadata": {},
   "source": [
    "### Invoking AgentCore Runtime with authorization\n",
    "\n",
    "Lets invoke the agent with the right authorization token type. In our case, it will be the Cognito access token. Copy the access token from the cell \"**Provision a Cognito User Pool**\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e07094ff-8f7b-41c7-aeae-6020bbc635e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#Update the Cognito access token here. Copy the access token from the cell \"Provision a Cognito User Pool\"\n",
    "cognito_bearer_token=\"<cognito-access-token>\"\n",
    "invoke_response = agentcore_runtime.invoke({\"prompt\": \"How is the weather now?\"}, bearer_token=cognito_bearer_token)\n",
    "invoke_response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c1d2bce-be41-478c-8bed-b4037c385795",
   "metadata": {},
   "source": [
    "### Invoking AgentCore Runtime with Python Http Client (Optional)\n",
    "\n",
    "Boto3 only supports IAM sigv4 based ingress. So for invoking your agent with bearer token you will need to use a regular python http client.  In the sample below we use request library in python."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4847ad03",
   "metadata": {},
   "source": [
    "#### Lets review the Agentcore configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6952cd2b",
   "metadata": {},
   "source": [
    "!cat .bedrock_agentcore.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f84e68d-6c04-41b9-bf5b-60edc3fa0985",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import urllib.parse\n",
    "import json\n",
    "import os\n",
    "\n",
    "from boto3.session import Session\n",
    "boto_session = Session()\n",
    "\n",
    "REGION_NAME = boto_session.region_name\n",
    "REGION_NAME\n",
    "\n",
    "BEDROCK_AGENT_CORE_ENDPOINT_URL=\"https://bedrock-agentcore.us-west-2.amazonaws.com\"\n",
    "\n",
    "#Update the agent ARN. You can get the agent arn from the above \"Review the AgentCore configuration\" section that prints the .bedrock_agentcore.yaml.\n",
    "invoke_agent_arn = \"<your-agent-arn>\"\n",
    "\n",
    "#Update the Cognito access token here. Copy the access token from the cell \"Provision a Cognito User Pool\"\n",
    "auth_token = \"<your-cognito-access-token>\"\n",
    "print(f\"Using Agent ARN from environment: {invoke_agent_arn}\")\n",
    "\n",
    "# URL encode the agent ARN\n",
    "escaped_agent_arn = urllib.parse.quote(invoke_agent_arn, safe='')\n",
    "\n",
    "# Construct the URL\n",
    "url = f\"{BEDROCK_AGENT_CORE_ENDPOINT_URL}/runtimes/{escaped_agent_arn}/invocations?qualifier=DEFAULT\"\n",
    "\n",
    "# Set up headers\n",
    "headers = {\n",
    "    \"Authorization\": f\"Bearer {auth_token}\",\n",
    "    \"X-Amzn-Trace-Id\": \"your-trace-id\", \n",
    "    \"Content-Type\": \"application/json\",\n",
    "    \"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id\": \"7a750a8c-11ab-447a-aec9-fe7b38402088222\"\n",
    "}\n",
    "\n",
    "# Enable verbose logging for requests\n",
    "import logging\n",
    "logging.basicConfig(level=logging.DEBUG)\n",
    "logging.getLogger(\"urllib3.connectionpool\").setLevel(logging.DEBUG)\n",
    "\n",
    "invoke_response = requests.post(\n",
    "    url,\n",
    "    headers=headers,\n",
    "    data=json.dumps({\"payload\": '{\"prompt\": \"Hello\"}'})\n",
    ")\n",
    "\n",
    "# Print response in a safe manner\n",
    "print(f\"Status Code: {invoke_response.status_code}\")\n",
    "print(f\"Response Headers: {dict(invoke_response.headers)}\")\n",
    "\n",
    "# Handle response based on status code\n",
    "if invoke_response.status_code == 200:\n",
    "    response_data = invoke_response.json()\n",
    "    print(\"Response JSON:\")\n",
    "    print(json.dumps(response_data, indent=2))  \n",
    "elif invoke_response.status_code >= 400:\n",
    "    print(f\"Error Response ({invoke_response.status_code}):\")\n",
    "    error_data = invoke_response.json()\n",
    "    print(json.dumps(error_data, indent=2))\n",
    "    \n",
    "else:\n",
    "    print(f\"Unexpected status code: {invoke_response.status_code}\")\n",
    "    print(\"Response text:\")\n",
    "    print(invoke_response.text[:500])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d3fdfe404469632",
   "metadata": {},
   "source": [
    "## Cleanup (Optional)\n",
    "\n",
    "Let's now clean up the AgentCore Runtime created"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76a6cf1416830a54",
   "metadata": {},
   "outputs": [],
   "source": [
    "from boto3.session import Session\n",
    "boto_session = Session()\n",
    "\n",
    "agentcore_control_client = boto3.client(\n",
    "    'bedrock-agentcore-control',\n",
    "    region_name=region\n",
    ")\n",
    "ecr_client = boto3.client(\n",
    "    'ecr',\n",
    "    region_name=region\n",
    "    \n",
    ")\n",
    "\n",
    "iam_client = boto3.client('iam')\n",
    "\n",
    "runtime_delete_response = agentcore_control_client.delete_agent_runtime(\n",
    "    agentRuntimeId=launch_result.agent_id,\n",
    "    \n",
    ")\n",
    "\n",
    "response = ecr_client.delete_repository(\n",
    "    repositoryName=launch_result.ecr_uri.split('/')[1],\n",
    "    force=True\n",
    ")\n",
    "\n",
    "policies = iam_client.list_role_policies(\n",
    "    RoleName=agentcore_iam_role['Role']['RoleName'],\n",
    "    MaxItems=100\n",
    ")\n",
    "\n",
    "for policy_name in policies['PolicyNames']:\n",
    "    iam_client.delete_role_policy(\n",
    "        RoleName=agentcore_iam_role['Role']['RoleName'],\n",
    "        PolicyName=policy_name\n",
    "    )\n",
    "iam_response = iam_client.delete_role(\n",
    "    RoleName=agentcore_iam_role['Role']['RoleName']\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b118ad38-feeb-4d1d-9d57-e5c845becc56",
   "metadata": {},
   "source": [
    "# Congratulations!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (Genesis)",
   "language": "python",
   "name": "genesis"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
